import numpy as np
import torch
import os
from seg_colors import segment_colors
import cv2

main_path = "../VOCdevkit/VOC2012"


def get_files():
    train_path = os.path.join(main_path, "ImageSets/Segmentation", "train.txt")
    val_path = os.path.join(main_path, "ImageSets/Segmentation", "val.txt")

    with open(train_path, "r", encoding="utf-8") as f:
        train_txt = f.readlines()

    with open(val_path, "r", encoding="utf-8") as f:
        val_txt = f.readlines()

    return train_txt, val_txt


def mat_resize(mat, dict_size=(512, 512)):
    h, w = mat.shape[:2]
    if h >= dict_size[0]:
        mat = mat[:dict_size[0], ...]
    else:
        mat = np.pad(mat, pad_width=((0, dict_size[0] - h), (0, 0), (0, 0)), mode="constant", constant_values=[0])
    if w > dict_size[1]:
        mat = mat[:, :dict_size[1], ...]
    else:
        mat = np.pad(mat, pad_width=((0, 0), (0, dict_size[1] - w), (0, 0)), mode="constant", constant_values=[0])

    return mat


def save_io_mat(file_list, name, mode="train"):
    seg_path = "SegmentationClass"
    origin_path = "JPEGImages"
    seg_color = segment_colors[name]["rgb"]
    inputs = []
    outputs = []
    for file_name in file_list:
        file_name = file_name.replace("\n", "")
        org_file_path = os.path.join(main_path, origin_path, file_name + ".jpg")
        seg_file_path = os.path.join(main_path, seg_path, file_name + ".png")
        org_mat = cv2.imread(org_file_path)
        seg_mat = cv2.imread(seg_file_path)
        if org_mat is None or seg_mat is None:
            continue
        org_mat = cv2.cvtColor(org_mat, cv2.COLOR_BGR2RGB)
        seg_mat = cv2.cvtColor(seg_mat, cv2.COLOR_BGR2RGB)

        org_mat = mat_resize(org_mat)
        seg_mat = mat_resize(seg_mat)

        seg_labels = np.zeros(seg_mat.shape[:2])
        seg_labels[np.all(seg_mat == seg_color, axis=-1)] = 1

        inputs.append(org_mat)
        outputs.append(seg_labels)
    inputs = np.stack(inputs, axis=0)
    outputs = np.stack(outputs, axis=0)
    features = torch.from_numpy(inputs).permute([0, 3, 1, 2])
    labels = torch.from_numpy(outputs)

    torch.save({"features": features, "labels": labels}, f"./data/voc12_{name}_{mode}.pt")


train_txt, val_txt = get_files()
save_io_mat(train_txt, "person",mode="train")
save_io_mat(val_txt, "person",mode="val")
